package com.brian.ratelimit.aop;

import com.brian.ratelimit.annotation.BrianLimit;
import com.google.common.util.concurrent.RateLimiter;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.Signature;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.io.PrintWriter;
import java.lang.reflect.Method;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

/**
 * @program: architect
 * @author: Brian Huang
 * @create: 2019-05-25 22
 *
 * 使用aop环绕通知方法，拦截所有springmvc请求，实现对接口服务实现限流
 * ###1.判断方法上是否有@BrianLimit注解
 * ###2.有该注解则使用反射技术获取到注解方法上的参数
 * ###3.调用原生的RetaLimiter方法创建令牌桶
 * ###4.如果获取令牌超时,调用服务降级方法
 * ###5.如果获取到令牌，调用实际进入的方法
 **/
@Component
@Aspect
public class RateLimitAop {

    private Map<String,RateLimiter> rateLimiterMap = new ConcurrentHashMap<>();


    @Pointcut("execution(public * com.brian.ratelimit.controller.*.*(..))")
    public void ptc(){}


    @Around(value = "ptc()")
    public Object around(ProceedingJoinPoint proceedingJoinPoint) throws Throwable {
        //1.判断方法上是否有@BrianLimit注解
        BrianLimit brianLimit = getAnnotationClass(proceedingJoinPoint);
        if(brianLimit == null){
           //没有注解则进入实际的方法中
            proceedingJoinPoint.proceed();
        }
        //2.有该注解则使用反射技术获取到注解方法上的参数
        //RateLimiter争对每一个url地址是单例的
        double permits = brianLimit.permitsPerSecond();
        long timeout = brianLimit.timeout();
        RateLimiter rateLimiter = null;
        String url = getUrl();//获取url
        //3.调用原生的RetaLimiter方法创建令牌桶
        if(rateLimiterMap.containsKey(url)){
            rateLimiter = rateLimiterMap.get(url);
        }else{
            rateLimiter = RateLimiter.create(permits);
            rateLimiterMap.put(url,rateLimiter);
        }

        //4.如果获取令牌超时,调用服务降级方法
        boolean acquire = rateLimiter.tryAcquire(timeout, TimeUnit.MILLISECONDS);
        if(!acquire){
            //走服务降级
            serviceDegradation();
            return null;
        }
        //5.如果获取到令牌，调用实际进入的方法
        return proceedingJoinPoint.proceed();
    }

    /**
     * 获取自定义注解类
     * @param proceedingJoinPoint
     * @return
     */
    private BrianLimit getAnnotationClass(ProceedingJoinPoint proceedingJoinPoint){
        MethodSignature signature = (MethodSignature) proceedingJoinPoint.getSignature();
        if(signature == null){
            return null;
        }
        Method method = signature.getMethod();
        BrianLimit declaredAnnotation = method.getDeclaredAnnotation(BrianLimit.class);
        return declaredAnnotation;
    }

    /**
     * 获取到request的url
     * @return
     */
    private String getUrl(){
        // 获取当前URL
        ServletRequestAttributes attributes = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
        HttpServletRequest request = attributes.getRequest();
        String requestURI = request.getRequestURI();
        return requestURI;
    }

    /**
     * 走服务降级
     */
    private void serviceDegradation() throws IOException {
        ServletRequestAttributes attributes = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
        HttpServletResponse response = attributes.getResponse();
        response.setHeader("Content-type", "text/html;charset=UTF-8");
        PrintWriter writer = response.getWriter();
        try {
            Map<String,String> map = new HashMap<>();
            map.put("code","200");
            map.put("message","亲,服务器忙！请稍后重试!");
            writer.println(map.toString());
        } catch (Exception e) {

        } finally {
            writer.close();
        }

    }
}
